# encoding=utf-8
import cv2
import numpy
import onnxruntime
import logging
import time
from PIL import Image

logging.basicConfig(level=logging.DEBUG)


class OnnxHelper:
    def __init__(self,
                 onnx_path: str,
                 preprocess,
                 postprocess):
        self.model = onnxruntime.InferenceSession(
            onnx_path, providers=onnxruntime.get_available_providers()
        )
        self.input_name = [ele.name for _, ele in enumerate(self.model.get_inputs())]
        self.output_name = [ele.name for _, ele in enumerate(self.model.get_outputs())]

        self.preprocess = preprocess
        self.postprocess = postprocess

    def infer(self, data_in):
        logging.info("Start to run, input_name: {}, output_name: {}".format(self.input_name, self.output_name))

        start = time.perf_counter()
        _s1 = time.perf_counter()
        preprocess_data = self.preprocess(data_in)
        end = time.perf_counter()
        logging.info("preprocess call time cost: {}s".format(end - start))

        input_pair = {}
        for name, data in zip(self.input_name, preprocess_data):
            input_pair[name] = data

        start = time.perf_counter()
        out_data = self.model.run(self.output_name, input_pair)
        end = time.perf_counter()
        logging.info("infer call time cost: {}s".format(end - start))

        start = time.perf_counter()
        post_data = self.postprocess(out_data)
        _s2 = time.perf_counter()
        end = time.perf_counter()
        logging.info("post call time cost: {}s".format(end - start))
        logging.info("all cost time: {}s".format(_s2 - _s1))
        return post_data


def pre(data_in: str):
    image = Image.open(data_in)
    _image = image.resize((512, 512), Image.BICUBIC)
    np_image = numpy.array(_image)
    _np_image = numpy.transpose(np_image, (2, 0, 1))
    __np_image = numpy.expand_dims(_np_image, axis=0)
    return [__np_image.astype(numpy.float32) / 255]


def post(data_in):
    out_np = data_in[0]
    y = numpy.squeeze(out_np) * 255.0

    # 下面照抄
    predict_list = []
    for index in range(y.shape[0]):
        predict = Image.fromarray(y[index].squeeze().astype(numpy.uint8), mode="L")  # 灰度模式存图，h * w无通道
        predict_resize = numpy.array(predict.resize((2816, 1880)))
        predict_list.append(predict_resize)

    # {"background": 0, "hemorrhages": 1, "hard_exudates": 2, "microaneurysms": 3, "disc": 4, "soft_exudates": 5}
    predict_list = predict_list[0], predict_list[2], predict_list[3], predict_list[1], predict_list[5], predict_list[4]
    prob_arr = numpy.array(predict_list) / 255.0
    prob_arr /= numpy.sum(prob_arr, axis=0)
    pred_arr = numpy.zeros_like(prob_arr).astype(numpy.uint8)
    pred_label_arr = numpy.argmax(prob_arr, axis=0)

    for j in range(pred_arr.shape[0]):
        pred_arr[j, :, :][pred_label_arr == j] = 255

    return pred_arr


def show(pred_arr: numpy.ndarray):
    obj_lst = ["background", "hemorrhages", "hard_exudates", "microaneurysms", "disc", "soft_exudates"]
    seg_color_list = [(0, 0, 0), (0, 0, 255), (0, 255, 0), (0, 255, 255), (255, 0, 0), (255, 0, 255)]

    pred_flag_arr = numpy.argmax(pred_arr, axis=0)
    pred_color_arr = numpy.zeros((pred_flag_arr.shape[0], pred_flag_arr.shape[1], 3), numpy.uint8)
    for i in range(1, len(obj_lst)):
        pred_color_arr[pred_flag_arr == i] = seg_color_list[i]

    cv2.imwrite("out.png", pred_color_arr)


if __name__ == '__main__':
    use_quant = True

    if use_quant:
        helper = OnnxHelper("unet-quant.onnx", pre, post)
    else:
        helper = OnnxHelper("unet-output.onnx", pre, post)

    for i in range(3):
        logging.info("{} call".format(i + 1))
        out = helper.infer("test_image.jpeg")
    show(out)
